from tree_sitter import Language, Parser, Node
from init_parser import *
from entity import If_node, TRUE_TO_FALSE
from util import merge_var_change, format_compare
from z3 import Real, Solver, sat, And, Or, Not, simplify, Int
import re
from copy import deepcopy
from typing import Dict, List  # 导入 Dict
from graphviz import Digraph
import os


# 2、ast中提取 if 语句结构，并处理变量的赋值，将结果储存在if_node
# 不是很好理解 还未理解
def format_tree(root: Node, if_tree: If_node, vars: Dict[str, str]):
    # 这个方法将一个语法树中的if语句抽象出来构造if树,只关注if语句,先忽略if语句块中的其他语句比如赋值和运算
    for child in root.children:
        # 这一步是查询出所有在if之前的对变量的修改
        # 判断当前节点是否只有一个子节点，如果是则可能是赋值语句
        if child.children.__len__() == 1:
            # 将当前节点的唯一子节点赋值给child1（赋值语句）
            child1 = child.children[0]
            # 检查当前节点是否为赋值语句
            if child1.type == 'assignment':
                # 赋值语句的左值（变量名）和右值（值/表达式）分别赋值
                l_value = str(child1.children[0].text, encoding='utf-8')
                r_value = str(child1.children[2].text, encoding='utf-8')
                # 如果赋值的右值是调用函数,说明赋值是通过函数调用完成的-不处理
                if child1.children[2].type == 'call':
                    # l_value暂时是它自身的符号
                    vars[l_value] = l_value
                else:
                    # 右值不是函数
                    # 如果l_value存在于词典，则之前该变量被赋值过，我们需要更新对照表
                    if l_value in vars:
                        vars[l_value] = vars[l_value].replace(l_value, f'({r_value})')
                    else:
                        # 如果l_value不在词典，则是第一次被赋值，变量名就是符号对照表对应的符号
                        vars[l_value] = r_value
        # 看当前节点是否为if语句节点
        if child.type == 'if_statement':
            # 获取if语句节点的所有子节点
            cc = child.children
            # 将一个新的if_node节点加入到if语法树
            if_tree.children.append(If_node('if', merge_var_change(str(cc[1].text, encoding='utf-8'), vars)))
            # 递归调用，下一步为if的block块，if_tree当前节点为刚刚加入的子节点
            # deepcopy(vars) 用于传递 vars 字典的副本，以避免在递归过程中修改原始字典
            format_tree(cc[3], if_tree.children[-1], deepcopy(vars))
            # 如果有else和elif，那么必定4项以后
            for c in cc[4:]:
                # 检查节点类型是否为elif子句
                if c.type == 'elif_clause':
                    if_tree.children.append(If_node(c.type.split('_')[0],
                                                    merge_var_change(str(c.children[1].text, encoding='utf-8'), vars)))
                    # 递归调用，下一步为elif的block块，if_tree当前节点为刚刚加入的子节点
                    format_tree(c.children[3], if_tree.children[-1], deepcopy(vars))
                # 检查节点类型是否为else子句
                if c.type == 'else_clause':
                    if_tree.children.append(If_node(c.type.split('_')[0], None))
                    # 递归调用，下一步为else的block块，if_tree当前节点为刚刚加入的子节点,else没有判断条件
                    format_tree(c.children[2], if_tree.children[-1], deepcopy(vars))


# 3、 生成路径覆盖测试用例的核心函数，为每个 if 节点生成所有可能的路径约束条件。它递归地处理每条路径，并生成相应的约束条件。
def get_truth_trees(root: If_node):
    '''
    使用路径覆盖
    '''

    # 考虑嵌套
    # 路径覆盖可以直接获得约束条件
    # 递归地生成从当前节点开始的路径上的约束条件
    def get_re(root: If_node, prefix: list, res: List[List[str]]):
        # 如果当前比较条件不是none
        if root.compare_statement != None:
            # 格式化条件并加入到prefix
            prefix.append(format_compare(root.compare_statement))
        i = root.children.__len__()
        # 当前节点没有子节点
        if i == 0:
            if prefix != []:
                # 将prefix加入到res
                res.append(prefix)
            return
        # 遍历当前节点的所有子节点
        for ii in range(i):
            # 从当前节点的第ii个子节点递归调用函数，生成从该节点开始的路径约束
            # deepcopy(prefix)是对prefix深拷贝
            get_re(root.children[ii], deepcopy(prefix), res)
            if root.children[ii].compare_statement != None:
                # 将当前子节点的条件的否定添加到prefix
                prefix.append(f'Not({format_compare(root.children[ii].compare_statement)})')

    l = len(root.children)
    res: List[List[str]] = []
    prefix = []  # 到达这个分支需要的前置条件
    for i in range(l):
        get_re(root.children[i], deepcopy(prefix), res)
        if root.children[i].compare_statement != None:
            prefix.append(f'Not({format_compare(root.children[i].compare_statement)})')
    return res
        # 如果当前比较条件不是none



# 4、打印 if 树结构，用于调试或可视化。它递归地遍历 If_node 树，并打印每个节点的详细信息
def print_tree(root: If_node, level=0):
    '''
    一个打印if树的方法
    '''
    # level表示当前节点在树中的层级  repr（root）返回root的字符串表示
    print('  ' * level + repr(root))
    for child in root.children:
        print_tree(child, level + 1)
        # module(None)
        # if ((a + 1) > 9)
        #     if (((a + 1) + 1) > 10)
        # elif ((b + 2) >= (a + 1) + 5)
        # else (None)


# 5、在字符串中替换变量名，根据给定的字典，将变量名替换为实际的值。
def replace_keys_with_values(string: str, rep_dict: dict):
    if string is None:
        return None
    for key, value in rep_dict.items():
        if key in string:
            string = string.replace(key, value)
            break
    return string


# 6、从字符串中提取变量名，用于提取约束条件中的变量，排除掉某些关键字。
def get_variables(s: str):
    variables = re.findall(r"\b[a-zA-Z]+\b", s)
    variables = list(filter(lambda x: x not in ['And', 'Not', 'Or'], variables))
    return variables


# 指定解析语言
PYTHON_LANGUAGE = Language('build/my-languages.so', 'python')
python_parser = Parser()
# parser的一个实例，设置解析语言为python
python_parser.set_language(PYTHON_LANGUAGE)


# 按钮-生成测试用例
# 1、入口，调用其他函数来解析代码、构建 if 树、生成约束条件，并最终输出测试用例。
def analyze(code: str):
    # ----------------------生成中间表示（抽象语法树）------------------------------
    result = ''
    # 解析python代码为抽象语法树，将字符串转换为byte，并将byte流输入到解析器，生成抽象语法树
    tree = python_parser.parse(bytes(code, "utf8"))
    # ----------------------进行符号执行-----------------------------------------
    # tree.root_node抽象语法树的根节点
    root_node = tree.root_node
    # If_node是自定义的类（在entity.py）中，表示if语句节点。下面是在构建if语法树的根节点
    if_tree = If_node('module', None)
    # 下面方法从抽象语法树根节点提取if语句相关结构，并将信息存到if语法树。
    # {}是空字典 作用是在符号执行遍历if语法树时更新符号变量对照表
    format_tree(root_node, if_tree, {})
    # ------------------------生成约束条件----------------------------------------
    # 打印if树
    # print_tree(if_tree)
    # 生成路径约束，一条路径的约束条件是一个数组
    restrictions = get_truth_trees(if_tree)
    # # 保存 if 树的图像
    # images_dir = 'static/images'
    # save_ast_image(if_tree, images_dir)
    # print(restrictions)
    # --------------------------构造Z3求解器求解约束--------------------------------
    letters = set()
    # 遍历约束条件列表
    for res in restrictions:
        # 遍历约束条件
        for r in res:
            # 从约束条件中提取变量名
            v = get_variables(r)
            # 变量名加到letters中
            letters.update(v)
    # print(letters)
    # 创建字典，包含Z3的逻辑操作符
    variables = {'And': And, 'Or': Or, 'Not': Not}
    # 遍历变量名
    for x in letters:
        # 为每一个变量x创建一个整型变量添加到字典
        variables[x] = Int(x)
    result = result + '一组路径覆盖测试用例如下\n'
    # 遍历约束列表
    for res in restrictions:
        # 针对一条路径的约束
        # 创建Z3求解器的实例
        s = Solver()
        # 针对当前约束列表中的一个约束
        for r in res:
            # eval（）将字符串形式的r转换为Z3表达式
            # simplify简化Z3表达式
            rr = simplify(eval(r, {}, variables))
            # print(rr)
            # 将简化后的约束条件添加到Z3求解器
            s.add(rr)
        result = result + f'该测试用例约束条件为:{res}\n'
        # 检查当前Z3求解器中的约束是否可满足，可满足返回sat
        if s.check() == sat:
            # 获取一个解
            m = s.model()
            result = result + str(m) + '\n'
        else:
            result = result + "无解\n"
    return result


if __name__ == "__main__":
    code = '''
    a = int(input())
    b = int(input())
    a = a + 1
    b = b + 2
    print(b)
    if a > 9:
        a = a + 1
        if a > 10:
            print('')
    elif b >= a + 5:
        print('22222')
    else:
        print('33333')
    '''
    print(analyze(code))


